#!/bin/bash

# ============================================================================
# vLLM Qwen3-8B 模型部署脚本
# ============================================================================
# 功能说明：
#   - 使用 vLLM 框架部署 Qwen3-8B 模型
#   - 提供 OpenAI 兼容的 API 服务
#   - 支持张量并行（Tensor Parallel）加速推理
#   - 限制 GPU 内存使用率，支持多任务共享 GPU
#
# 使用方法：
#   1. 确保已激活正确的 conda 环境: conda activate ragen
#   2. 确保已安装 vllm: pip install vllm
#   3. 运行脚本: bash deploy_vllm_qwen3.sh
#   4. 服务启动后，可通过 http://localhost:8177 访问 API
#
# API 测试示例：
# curl http://localhost:8177/v1/completions \
# -H "Content-Type: application/json" \
# -d '{
#     "model": "/home/yangcx24/Jayx/Models/Qwen3-8B",
#     "prompt": "你好，请介绍一下你自己。",
#     "max_tokens": 100,
#     "temperature": 0.7
# }'
# ============================================================================

# ============================================================================
# 配置参数
# ============================================================================

# 指定使用的 GPU（仅使用 GPU 0）
export CUDA_VISIBLE_DEVICES=0,1,2,3

# 模型路径（从 config/base.yaml 中获取）
MODEL_PATH="/home/yangcx24/Jayx/Models/Qwen3-32B"
SERVED_MODEL_NAME="Qwen3-32B"
# 张量并行度（使用 1 路张量并行，单 GPU 推理）
# 注意：单 GPU 时设置为 1
TENSOR_PARALLEL_SIZE=4

# GPU 内存使用率限制（0.0-1.0 之间，0.3 表示仅使用 30%）
# 设置较低的值可以让其他任务共享 GPU 资源
GPU_MEMORY_UTILIZATION=0.95

# 服务端口号
PORT=8177

# 服务主机地址（0.0.0.0 表示接受所有网络接口的连接）
HOST="0.0.0.0"

# 最大模型序列长度（可选，根据需要调整）
# Qwen3-8B 的 max_position_embeddings=40960（从模型 config.json 获取）
# 这是模型真实支持的最大上下文长度，不应超过此值
# 40960 tokens = 输入约10k + 输出30k，足够支持 api.yaml 的配置
MAX_MODEL_LEN=40960

# 其他可选参数
# 是否启用分词器的信任远程代码（Qwen 模型通常需要）
TRUST_REMOTE_CODE="--trust-remote-code"

# 是否启用日志详细模式
# VERBOSE="--verbose"

# ============================================================================
# 环境检查
# ============================================================================

echo "============================================"
echo "vLLM Qwen3-8B 模型部署脚本"
echo "============================================"
echo ""

# 检查 conda 环境
if [ -z "$CONDA_DEFAULT_ENV" ]; then
    echo "警告: 未检测到活动的 conda 环境"
    echo "建议运行: conda activate ragen"
    echo ""
fi

# 检查模型路径是否存在
if [ ! -d "$MODEL_PATH" ]; then
    echo "错误: 模型路径不存在: $MODEL_PATH"
    echo "请检查模型是否已下载到指定位置"
    exit 1
fi

echo "✓ 模型路径: $MODEL_PATH"

# 检查 vllm 是否已安装
if ! command -v vllm &> /dev/null; then
    echo "错误: vllm 未安装"
    echo "请运行: pip install vllm"
    exit 1
fi

echo "✓ vLLM 已安装"

# 检查可用的 GPU 数量
if command -v nvidia-smi &> /dev/null; then
    GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l)
    echo "✓ 检测到 $GPU_COUNT 张 GPU"
    
    if [ $GPU_COUNT -lt $TENSOR_PARALLEL_SIZE ]; then
        echo "警告: 可用 GPU 数量 ($GPU_COUNT) 少于张量并行度 ($TENSOR_PARALLEL_SIZE)"
        echo "建议将 TENSOR_PARALLEL_SIZE 设置为不大于 $GPU_COUNT"
    fi
else
    echo "警告: 无法检测 GPU，请确保 NVIDIA 驱动已正确安装"
fi

echo ""

# ============================================================================
# 启动 vLLM 服务
# ============================================================================

echo "正在启动 vLLM 服务..."
echo "配置信息:"
echo "  - 模型: $MODEL_PATH"
echo "  - 张量并行度: $TENSOR_PARALLEL_SIZE"
echo "  - GPU 内存使用率: $(echo "$GPU_MEMORY_UTILIZATION * 100" | bc)%"
echo "  - 服务地址: http://$HOST:$PORT"
echo "  - 最大序列长度: $MAX_MODEL_LEN tokens (模型真实最大长度)"
echo "  - 支持配置: 输入约 10k tokens + 输出 30k tokens = 40k tokens"
echo ""
echo "按 Ctrl+C 停止服务"
echo "============================================"
echo ""

# 构建 vllm serve 命令
# 使用 vllm serve 命令启动 OpenAI 兼容的 API 服务器
vllm serve "$MODEL_PATH" \
    --host "$HOST" \
    --port "$PORT" \
    --tensor-parallel-size "$TENSOR_PARALLEL_SIZE" \
    --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
    --max-model-len "$MAX_MODEL_LEN" \
    --served-model-name "$SERVED_MODEL_NAME" \
    --no-enable-prefix-caching
    $TRUST_REMOTE_CODE \
    # --enable-prefix-caching \  # 可选: 启用前缀缓存以提高性能
    # --disable-log-requests \   # 可选: 禁用请求日志
    # $VERBOSE                   # 可选: 启用详细日志

# ============================================================================
# 说明
# ============================================================================
# 
# 主要参数说明：
# 
# --tensor-parallel-size: 
#   张量并行度，将模型分片到多张 GPU 上进行并行推理
#   建议设置为 2 的幂次（1, 2, 4, 8 等）
#   值越大推理越快，但需要更多 GPU
#
# --gpu-memory-utilization:
#   GPU 内存使用率限制（0.0-1.0）
#   较低的值可以让其他任务共享 GPU，但可能降低吞吐量
#   推荐值: 0.3-0.5（多任务共享）, 0.8-0.95（单任务专用）
#
# --max-model-len:
#   最大序列长度（tokens），包括输入和输出
#   较大的值支持更长的上下文，但需要更多显存
#   如果显存不足，可以适当降低此值
#
# --trust-remote-code:
#   允许执行模型仓库中的远程代码
#   Qwen 系列模型通常需要此选项
#
# --enable-prefix-caching:
#   启用前缀缓存，可以提高相似请求的处理速度
#   适用于有大量相同前缀的场景（如系统提示词）
#
# --host 和 --port:
#   服务监听的地址和端口
#   0.0.0.0 表示接受所有网络接口的连接
#   如果只在本机访问，可以设置为 127.0.0.1
#
# ============================================================================
# API 端点说明：
# 
# vLLM 提供 OpenAI 兼容的 API 接口，主要端点包括：
# 
# 1. Completions API (文本补全):
#    POST http://localhost:8000/v1/completions
# 
# 2. Chat Completions API (对话):
#    POST http://localhost:8000/v1/chat/completions
# 
# 3. Models API (查询可用模型):
#    GET http://localhost:8000/v1/models
# 
# 详细 API 文档请参考: https://docs.vllm.ai/en/latest/serving/openai_compatible_server.html
# 
# ============================================================================

